from geometry.layer import Layer
import geopandas as gpd
import pandas as pd
import numpy as np
from geometry.geometry import Point, LineString, Polygon
from io import StringIO


class GeometryConverter:
    def __init__(self):
        self.conn = None

    def setDBConnection(self, conn):
        self.conn = conn

    def csvToLayer(self, csv_path):
        df = pd.read_csv(csv_path)
        for idx in df.index:
            pt = Point()
            pt.fromXY(df.loc[idx, "x"], df.loc[idx, "y"])
            df.loc[idx, "my_geom"] = pt

        L = Layer()
        L.setData(df)
        L.setGeometryType("Point")

        x_all = df["my_geom"].apply(lambda x: x.x)
        y_all = df["my_geom"].apply(lambda x: x.y)
        minx = min(x_all)
        miny = min(y_all)
        maxx = max(x_all)
        maxy = max(y_all)
        L.setBBox(minx, maxx, miny, maxy)
        L.name = csv_path.split("/")[-1].split(".")[0]
        return L

    def shpToLayer(self, shp_path):
        """
        1. 使用geopandas读取shp文件，得到一个geodataframe。此时几何图形在geometry列中，为shapely对象。提示：pandas基础操作见网盘里machine learning压缩包里面的教程，或者自行百度
        2. 判断几何图形的类型。提示：使用几何对象.geometryType()方法  #K 使用了geopandas 里面的.geom_type参数
        3. 把shapely几何图形转换为我们的几何图形对象。提示：list(shapely几何对象.coords)获取shapely几何对象的坐标列表，用fromXY将坐标列表导入为我们自定义的几何对象
        计算layer的范围，使用layer.setBBox()存入layer。bbox格式：XMin, XMax, YMin, YMax
        4. geodataframe转换为普通dataframe。删除原来的geometry列，加入我们的几何图形对象列，列名为my_geom
        5. 调用Layer.setData(dataframe)，将dataframe存入图层对象。返回这个图层对象

        :param shp_path:
        :return:
        """
        data = gpd.read_file(shp_path)
        """判断几何类型，并导入自定义的几何类型中"""
        newcol = []
        if data.shape[0] == 0:
            geomType = None
        else:
            geomType = data.loc[0, 'geometry'].geom_type
        for g in data['geometry']:
            # print(g.geom_type)
            if g.geom_type.lower() == "polygon":
                p = Polygon()
                p.fromText(str(g))
                newcol.append(p)
            elif g.geom_type.lower() == "point":
                po = Point()
                po.fromText(str(g))
                newcol.append(po)
            elif g.geom_type.lower() == "linestring":
                ls = LineString()
                ls.fromText(str(g))
                newcol.append(ls)
            else:
                print(g.geom_type, "暂未定义的数据类型………不支持导入")
        # 将geodataframe转换成dataframe,并替换掉‘geometry’列

        df1 = pd.DataFrame(data.drop(columns='geometry'))
        df1['my_geom'] = newcol  # 需要传入20个元素
        # 调用Layer.setData(dataframe)，将dataframe存入图层对象。返回这个图层对象
        L = Layer()
        L.setData(df1)
        L.setGeometryType(geomType)
        """输出显示"""
        # 取出shp文件中的minx,miny,maxx,maxy;并传入layer中
        minx = min(data['geometry'].bounds['minx'])
        miny = min(data['geometry'].bounds['miny'])
        maxx = max(data['geometry'].bounds['maxx'])
        maxy = max(data['geometry'].bounds['maxy'])
        L.setBBox(minx, maxx, miny, maxy)
        L.name = shp_path.split("/")[-1].split(".")[0]
        return L

    def layerToShp(self, layer, save_path):
        """
        1. 我们的几何对象->shapely的几何对象
        2. 创建包含属性、shapely几何对象的GeoDataFrame
        2. gdf.to_file()

        :param layer:
        :param save_path:
        :return:
        """

        if save_path[-4:] != ".shp":
            save_path += ".shp"
        data = layer.getdata()
        print(data)
        print(data.columns)
        data_wkt = []
        for l in data['my_geom']:  # 可以考虑使用apply替代for循环，提高速度：data_wkt = data["my_geom"].apply(lambda x: x.asText())
            data_wkt.append(l.asText())
        dg = gpd.GeoSeries.from_wkt(data_wkt)
        gdf = gpd.GeoDataFrame(data.drop('my_geom', axis=1), geometry=dg)
        gdf.to_file(save_path, encoding='utf-8')  # 避免非英文字符乱码

    def postgisToLayer(self, table_name):
        # 使用pandas将数据表读取为一个dataframe。几何数据的处理：在sql语句里使用st_astext(geom)获取wkt
        df = pd.read_sql(f"select *, st_astext(geom) as my_wkt from {table_name}", self.conn)
        # pandas读取的geom列是wkb(well known binary)，删掉geom列
        df = df.drop('geom', 1)  # drop删除一列，geom是列名，1代表列（0代表行）

        newcol = []  # df的新列，存放我们的几何对象

        geomType = pd.read_sql(f"select type from geometry_columns where f_table_name='{table_name}' ", self.conn).loc[0, "type"]
        # 然后将wkt转换为我们的几何对象，放入df的my_geom列
        if geomType == "LINESTRING":
            geomType = "LineString"
            for g in df['my_wkt']:
                ls = LineString()
                ls.fromText(str(g))
                newcol.append(ls)
        elif geomType == "POINT":
            geomType = "Point"
            for g in df['my_wkt']:
                po = Point()
                po.fromText(str(g))
                newcol.append(po)
        elif geomType == 'POLYGON':
            geomType = "Polygon"
            for g in df['my_wkt']:
                p = Polygon()
                p.fromText(str(g))
                newcol.append(p)
        else:  # 不符合我们的几何类型，应该报错（暂未实现）
            # print(g.geom_type, "暂未定义的数据类型………不支持导入")
            pass

        df['my_geom'] = newcol

        L = Layer()  # 创建一个图层对象layer
        L.setData(df)  # 调用layer.setData()，把dataframe存入layer
        L.setGeometryType(geomType)  # 给layer.geometryType属性赋值
        L.name = table_name  # 图层名 = 表名

        # 计算layer的范围，使用layer.setBBox()存入layer。bbox格式：XMin, XMax, YMin, YMax。使用geom.bbox获取LineString和Polygon的bbox。Point使用X()和Y()获取坐标

        df_box = pd.read_sql(f"select "
                             f"st_xmin(st_extent(geom)) as xmin, "
                             f"st_xmax(st_extent(geom)) as xmax, "
                             f"st_ymin(st_extent(geom)) as ymin, "
                             f"st_ymax(st_extent(geom)) as ymax "
                             f" from {table_name}", self.conn)
        xmin = df_box.at[0, 'xmin']
        xmax = df_box.at[0, 'xmax']
        ymin = df_box.at[0, 'ymin']
        ymax = df_box.at[0, 'ymax']
        if xmin is None:
            xmin = np.inf
        if xmax is None:
            xmax = -np.inf
        if ymin is None:
            ymin = np.inf
        if ymax is None:
            ymax = -np.inf
        L.setBBox(xmin, xmax, ymin, ymax)

        return L  # 返回图层对象

    def layerToPostgis(self, layer, save_name):
        # 判断表是否存在
        count = pd.read_sql(f"select count(*) as count from pg_class where relname = '{save_name}'", self.conn).loc[
            0, "count"]
        assert count == 0, f"表{save_name}已存在"

        # 在数据库里创建表。几何图形字段的名称统一设置为geom
        df = layer.getdata().copy()
        df["geom"] = df["my_geom"].apply(lambda x: x.asText())
        for drop_col in ["my_geom", "my_wkt"]:
            if drop_col in df.columns:
                df.drop(columns=drop_col, inplace=True)
        col_list = []
        for col in df.columns:
            if col == "geom":
                dtype = f"geometry({layer.geometryType}, 0)"
            else:
                dtype = "text"
            col_list.append(col.lower() + " " + dtype)
        with self.conn:
            with self.conn.cursor() as cursor:
                sql = f"""
                    create table {save_name}(
                        {",".join(col_list)}
                    )
                """
                cursor.execute(sql)

        # 导入数据
        with self.conn:
            with self.conn.cursor() as cur:
                columns = list(df.columns)
                output = StringIO()
                df.to_csv(output, sep='\t', index=False, header=False)
                output1 = output.getvalue()
                cur.copy_from(StringIO(output1), save_name, columns=columns)

